"""火山方舟模型类（支持多模态和文本对话）"""
import requests
import json
import logging
from typing import Dict, List, Any

class VolcanoArkModel:
    """火山方舟模型实现类（支持图文分析和文本对话）"""
    
    API_URL = "https://ark.cn-beijing.volces.com/api/v3/chat/completions"
    
    def __init__(self, api_key: str, model: str = "doubao-seed-1-6-thinking-250615"):
        """
        初始化模型参数
        
        :param api_key: 火山方舟API密钥（从控制台获取）
        :param model: 模型ID（默认为doubao-seed-1-6-thinking-250615）
        """
        self.api_key = api_key
        self.model = model
        self.logger = logging.getLogger("VolcanoArkModel")
        self._setup_logger()
    
    def _setup_logger(self):
        """配置日志记录器"""
        self.logger.setLevel(logging.INFO)
        handler = logging.StreamHandler()
        formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
        handler.setFormatter(formatter)
        self.logger.addHandler(handler)
        self.logger.info("火山方舟模型初始化完成 | 模型: %s", self.model)
    
    def send_request(self, messages: List[Dict], **kwargs) -> Dict[str, Any]:
        """
        发送通用API请求
        
        :param messages: 消息列表
        :param kwargs: 额外参数（max_tokens, temperature等）
        :return: API响应字典
        """
        # 构建请求体
        payload = {
            "model": self.model,
            "messages": messages
        }
        
        # 添加可选参数
        optional_params = ["max_tokens", "temperature", "top_p", "top_k", 
                          "repetition_penalty", "stop", "stream"]
        for param in optional_params:
            if param in kwargs:
                payload[param] = kwargs[param]
        
        # 配置请求头
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self.api_key}"
        }
        
        try:
            # 发送请求（超时设置30秒）
            response = requests.post(
                self.API_URL,
                headers=headers,
                json=payload,
                timeout=30
            )
            
            # 处理HTTP错误
            if response.status_code != 200:
                error_msg = f"API错误: {response.status_code} - {response.text}"
                self.logger.error(error_msg)
                return {"error": error_msg}
            
            return response.json()
        
        except requests.exceptions.RequestException as e:
            error_msg = f"网络请求异常: {str(e)}"
            self.logger.error(error_msg)
            return {"error": error_msg}

    def text_chat(self, prompt: str, **kwargs) -> str:
        """
        执行纯文本对话
        
        :param prompt: 文本提示
        :param kwargs: 可选参数（max_tokens, temperature等）
        :return: 模型生成的文本响应
        """
        messages = [
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": prompt}
                ]
            }
        ]
        
        response = self.send_request(messages, **kwargs)
        return self._parse_response(response)
    
    def image_with_text(self, text_prompt: str, image_url: str) -> str:
        """
        执行图文分析
        
        :param text_prompt: 文本提示（如"图片主要讲了什么?"）
        :param image_url: 公网可访问的图片URL
        :return: 模型生成的文本响应
        """
        messages = [
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": text_prompt},
                    {"type": "image_url", "image_url": {"url": image_url}}
                ]
            }
        ]
        
        response = self.send_request(messages)
        return self._parse_response(response)
    
    def _parse_response(self, response_data: Dict[str, Any]) -> str:
        """解析火山方舟API响应"""
        if "error" in response_data:
            return response_data["error"]
        
        try:
            if "choices" in response_data and response_data["choices"]:
                first_choice = response_data["choices"][0]
                message = first_choice.get("message")
                
                if message and "content" in message:
                    return message["content"].strip()
            
            # 记录异常响应
            self.logger.warning("响应格式异常: %s", json.dumps(response_data, indent=2))
            return "火山方舟返回了无法解析的响应格式"
        
        except Exception as e:
            self.logger.exception("响应解析异常: %s", str(e))
            return f"响应解析失败: {str(e)}"

# 使用示例
if __name__ == "__main__":
    # 初始化模型（替换为您的真实API密钥）
    model = VolcanoArkModel(
        api_key="9199f23c-32b1-45b3-9196-7c33da9f5c12",
        model="doubao-seed-1-6-thinking-250615"
    )
    
    # 测试图文分析
    image_response = model.image_with_text(
        text_prompt="图片主要讲了什么?",
        image_url="https://ark-project.tos-cn-beijing.ivolces.com/images/view.jpeg"
    )
    
    print("图文分析结果:")
    print(image_response)
    
    # 测试文本对话
    text_response = model.text_chat(
        "2025年中国AI行业将如何发展？请用100字以内回答",
        max_tokens=200,
        temperature=0.7
    )
    
    print("\n文本对话结果:")
    print(text_response)

